import numpy as np
import os
import logging, json
import fiftyone.utils.data as foud
import fiftyone.core.metadata as fom
import fiftyone.core.labels as fol
import tqdm


class HPCompetitionLabeledImageDatasetImporter(foud.LabeledImageDatasetImporter):
    """Custom importer for labeled image datasets.

    Args:
        dataset_dir (None): the dataset directory. This may be optional for
            some importers
        shuffle (False): whether to randomly shuffle the order in which the
            samples are imported
        seed (None): a random seed to use when shuffling
        max_samples (None): a maximum number of samples to import. By default,
            all samples are imported
        **kwargs: additional keyword arguments for your importer
    """

    def __init__(
        self,
        dataset_dir=None,
        shuffle=False,
        seed=None,
        max_samples=None,
        imgs_dir=None,
        threshold=0.01,
        field_name="predict",
        **kwargs,
    ):
        super().__init__(
            dataset_dir=dataset_dir,
            shuffle=shuffle,
            seed=seed,
            max_samples=max_samples,
        )
        self.imgs_dir = imgs_dir
        assert imgs_dir is not None, "imgs_dir should not be None, it should point to the root of test datset "
        self.threshold = threshold
        self.field_name = field_name

    def __len__(self):
        """The total number of samples that will be imported.

        Raises:
            TypeError: if the total number is not known
        """
        # Return the total number of samples in the dataset (if known)
        return len(self.samples)

    def __next__(self):
        """Returns information about the next sample in the dataset.

        Returns:
            an  ``(image_path, image_metadata, label)`` tuple, where

            -   ``image_path``: the path to the image on disk
            -   ``image_metadata``: an
                :class:`fiftyone.core.metadata.ImageMetadata` instances for the
                image, or ``None`` if :meth:`has_image_metadata` is ``False``
            -   ``label``: an instance of :meth:`label_cls`, or a dictionary
                mapping field names to :class:`fiftyone.core.labels.Label`
                instances, or ``None`` if the sample is unlabeled

        Raises:
            StopIteration: if there are no more samples to import
        """
        # Implement loading the next sample in your dataset here
        image_path, image_metadata, detections = next(self.image_iter)
        # print(image_metadata)
        return image_path, image_metadata, {self.field_name: detections}

    @property
    def has_dataset_info(self):
        """Whether this importer produces a dataset info dictionary."""
        return True

    @property
    def has_image_metadata(self):
        """Whether this importer produces
        :class:`fiftyone.core.metadata.ImageMetadata` instances for each image.
        """
        return True

    @property
    def label_cls(self):
        """The :class:`fiftyone.core.labels.Label` class(es) returned by this
        importer.

        This can be any of the following:

        -   a :class:`fiftyone.core.labels.Label` class. In this case, the
            importer is guaranteed to return labels of this type
        -   a list or tuple of :class:`fiftyone.core.labels.Label` classes. In
            this case, the importer can produce a single label field of any of
            these types
        -   a dict mapping keys to :class:`fiftyone.core.labels.Label` classes.
            In this case, the importer will return label dictionaries with keys
            and value-types specified by this dictionary. Not all keys need be
            present in the imported labels
        -   ``None``. In this case, the importer makes no guarantees about the
            labels that it may return
        """
        # Return the appropriate value here
        return {self.field_name: fol.Detections}
    def close(self, *args):
        if self.tmp_dataset_dir:
            os.system(f"rm -rf {self.tmp_dataset_dir}")

    
    def setup(self):
        """Performs any necessary setup before importing the first sample in
        the dataset.

        This method is called when the importer's context manager interface is
        entered, :func:`DatasetImporter.__enter__`.
        """
        if os.path.isfile(self.dataset_dir) and self.dataset_dir.endswith(".zip"):
            self.zip_file = self.dataset_dir
            self.dataset_dir = 'tmpxxxxxxxllxxxxlddd'
            os.makedirs(self.dataset_dir, exist_ok=False)
            os.system(f"unzip {self.zip_file} -d {self.dataset_dir}")
            self.tmp_dataset_dir = self.dataset_dir
            
        if not os.path.isdir(self.dataset_dir):
            raise Exception("Invalid sumbit directory")
        self.samples = []
        subset_json_files = os.listdir(self.dataset_dir)
        for json_file in subset_json_files:
            subset_name = json_file.split(".")[0]
            if not json_file.endswith(".json"):
                continue
            print(f"loading {subset_name}.json")
            with open(os.path.join(self.dataset_dir, json_file), "r") as f:
                datas = json.load(f)
                detections = {}
                image_metadatas = {}
                for data in tqdm.tqdm(datas):
                    img_path = data['name']
                    img_path = os.path.join(self.imgs_dir, subset_name, "Images", img_path)
                    # tag = img_path.split('/')[-3]
                    image_metadatas[img_path] = fom.ImageMetadata.build_for(os.path.join(self.imgs_dir, subset_name, "Images", img_path))
                    width, height = image_metadatas[img_path].width, image_metadatas[img_path].height
                    bbox = data['bbox']
                    score = data['score']
                    label = str(data['category'])
                    if score < self.threshold:
                        continue
                    x1, y1, x2, y2 = bbox
                    x1, y1, x2, y2 = float(x1)/width, float(y1)/height, float(x2)/width, float(y2)/height
                    if img_path not in detections:
                        detections[img_path] = []
                    detections[img_path].append(fol.Detection(label=label, bounding_box=[x1, y1, x2-x1, y2-y1], confidence=score))
                for img_path, detections in detections.items():
                    fo_detections = fol.Detections(detections=detections)
                    self.samples.append((img_path, image_metadatas[img_path], fo_detections))
            print(f'loaded {len(self.samples)} samples')
            
            
        self.image_iter = iter(self.samples)
        
        self.info = {"dataset_dir": self.dataset_dir, }

    def get_dataset_info(self):
        """Returns the dataset info for the dataset.

        By convention, this method should be called after all samples in the
        dataset have been imported.

        Returns:
            a dict of dataset info
        """
        return self.info
    
class HPCompetitionLabeledImageDatasetExporter(foud.LabeledImageDatasetExporter):
    """Custom exporter for labeled image datasets.

    Args:
        export_dir (None): the directory to write the export. This may be
            optional for some exporters
        **kwargs: additional keyword arguments for your exporter
    """

    def __init__(self, export_dir=None, **kwargs):
        super().__init__(export_dir=export_dir)

    @property
    def requires_image_metadata(self):
        """Whether this exporter requires
        :class:`fiftyone.core.metadata.ImageMetadata` instances for each sample
        being exported.
        """
        return True

    @property
    def label_cls(self):
        """The :class:`fiftyone.core.labels.Label` class(es) exported by this
        exporter.

        This can be any of the following:

        -   a :class:`fiftyone.core.labels.Label` class. In this case, the
            exporter directly exports labels of this type
        -   a list or tuple of :class:`fiftyone.core.labels.Label` classes. In
            this case, the exporter can export a single label field of any of
            these types
        -   a dict mapping keys to :class:`fiftyone.core.labels.Label` classes.
            In this case, the exporter can handle label dictionaries with
            value-types specified by this dictionary. Not all keys need be
            present in the exported label dicts
        -   ``None``. In this case, the exporter makes no guarantees about the
            labels that it can export
        """
        return fol.Detections

    def setup(self):
        """Performs any necessary setup before exporting the first sample in
        the dataset.

        This method is called when the exporter's context manager interface is
        entered, :func:`DatasetExporter.__enter__`.
        """
        # Your custom setup here
        self.samples = {}

    def log_collection(self, sample_collection):
        """Logs any relevant information about the
        :class:`fiftyone.core.collections.SampleCollection` whose samples will
        be exported.

        Subclasses can optionally implement this method if their export format
        can record information such as the
        :meth:`fiftyone.core.collections.SampleCollection.name` and
        :meth:`fiftyone.core.collections.SampleCollection.info` of the
        collection being exported.

        By convention, this method must be optional; i.e., if it is not called
        before the first call to :meth:`export_sample`, then the exporter must
        make do without any information about the
        :class:`fiftyone.core.collections.SampleCollection` (which may not be
        available, for example, if the samples being exported are not stored in
        a collection).

        Args:
            sample_collection: the
                :class:`fiftyone.core.collections.SampleCollection` whose
                samples will be exported
        """
        pass
    
    def convert_onetxt2json(self, detections: fol.Detections, fileName: str, metadata: dict):
        result_json =[]
        width = metadata.width
        height = metadata.height
        for detection in detections.detections:
            label = detection.label
            score = detection.confidence
            xywh = detection.bounding_box
            xyxy = [xywh[0]*width, xywh[1]*height, (xywh[0]+xywh[2])*width, (xywh[1]+xywh[3])*height]

            score = max(0, min(1.0, score))
            one_json = {
                "name" : os.path.split(fileName)[1],
                "category" : int(label),
                "bbox":  xyxy,
                "score": float(score)
            }
            result_json.append(one_json)
        return result_json
    
    def export_sample(self, image_or_path, label, metadata=None):
        """Exports the given sample to the dataset.

        Args:
            image_or_path: an image or the path to the image on disk
            label: an instance of :meth:`label_cls`, or a dictionary mapping
                field names to :class:`fiftyone.core.labels.Label` instances,
                or ``None`` if the sample is unlabeled
            metadata (None): a :class:`fiftyone.core.metadata.ImageMetadata`
                instance for the sample. Only required when
                :meth:`requires_image_metadata` is ``True``
        """
        sample = self.convert_onetxt2json(label, image_or_path, metadata)
        subdir = image_or_path.split("/")[-3]
        if subdir not in self.samples:
            self.samples[subdir] = []
        self.samples[subdir].extend(sample)
        

    def close(self, *args):
        """Performs any necessary actions after the last sample has been
        exported.

        This method is called when the importer's context manager interface is
        exited, :func:`DatasetExporter.__exit__`.

        Args:
            *args: the arguments to :func:`DatasetExporter.__exit__`
        """
        os.makedirs(self.export_dir, exist_ok=True)
        for subdir, samples in self.samples.items():
            with open(os.path.join(self.export_dir, subdir+".json"), "w") as tf:
                json.dump(samples, tf, indent=4, separators=(',',': '))
        if os.path.isfile(os.path.join(self.export_dir, "submit.zip")):
            os.remove(os.path.join(self.export_dir, "submit.zip"))
        os.system(f"cd {self.export_dir} && zip -r submit.zip *.json")
        self.samples = {}


import fiftyone.types as fot

class HPCompetitionImageDataset(fot.LabeledImageDataset):
    def get_dataset_importer_cls(self):
        return HPCompetitionLabeledImageDatasetImporter
    
    def get_dataset_exporter_cls(self):
        return HPCompetitionLabeledImageDatasetExporter
        